#pragma once
#include <iostream>
#include <sys/io.h>
#include <fstream>
#include <filesystem>


#include <torch/torch.h>
#include <c10/core/TensorImpl.h>
#include <torch/csrc/autograd/grad_mode.h>
#include <torch/script.h>



#include <opencv2/core.hpp>  
#include <opencv2/highgui.hpp>  
#include <opencv2/imgproc.hpp>
#include <opencv2/dnn.hpp>
#include <time.h>  
#include <string>
#include <chrono>
#include <time.h>

enum YOLO_Status {
	YOLO_NORMAL,
	YOLO_MODELNOTAVAILABLE,
	YOLO_CANNOTFORWARD,
	YOLO_CLASSNAMESLOST
};


struct Net_config
{
	float confThreshold; // Confidence threshold
	float iouThreshold;  // Non-maximum suppression threshold
	int batchSize = 1;
	bool useGPU = false;
	int gpuId = 0;
	int imgWidth = 640;
	int imgHeight = 640;
	unsigned long long max_gpu_memory = 2147483648;
	std::string classesFile;
	std::string modelPath;
};

enum Det
{
	tl_x = 0,
	tl_y = 1,
	br_x = 2,
	br_y = 3,
	score = 4,
	class_idx = 5
};


struct Detection
{
	cv::Rect bbox;
	float score;
	int class_idx;
	std::string className;
	int64_t times = 0;
};

struct FrameInfo {
	int64_t times=-1;
	cv::Mat frame;
};


class YOLOV7
{
public:
	YOLOV7(Net_config config);
	YOLOV7(std::string classesFile, std::string modelPath, 
		int batchSize=1,int imgWidth=640,int imgHeight=640,
		float confThreshold = (double)0.25, float iouThreshold = (double)0.45, bool useGPU = false, int gpuId = 0, long max_gpu_memory = 2147483648);
	~YOLOV7();
	std::vector<Detection> detect(cv::Mat& frame);
	bool addBatchItem(cv::Mat& frame);
	bool canAddItem();
	int hasItems();
	void setBatchSize(int batchSize=1);
	std::vector<std::vector<Detection>> detectBatch();
	static cv::Mat drawResult(cv::Mat& img,const std::vector<Detection> detections,bool label = true);
	YOLO_Status getStatus();
private:
	YOLO_Status status = YOLO_NORMAL;
	int gpuId = 0;
	std::string errorCode;
	int inpWidth;
	int inpHeight;
	int nout;
	int num_proposal;
	int batchSize;
	std::vector<std::string> class_names;
	int num_class;
	bool cudaAvailable = false;
	torch::DeviceType device_type;
	torch::Device *device;
	torch::jit::script::Module module;
	std::string classesFile;
	std::string modelPath;


	float confThreshold;
	float iouThreshold;
	std::vector<torch::Tensor> tensorList;
	std::vector<std::vector<float>> batchPadInfo;
	std::vector<std::vector<int>> batchImageSize;

	void initialize(Net_config config);
	std::vector<std::string> LoadNames(const std::string& path);
	std::vector<float> LetterboxImage(const cv::Mat& src, cv::Mat& dst, const cv::Size& out_size);
	void Tensor2Detection(const at::TensorAccessor<float, 2>& offset_boxes, const at::TensorAccessor<float, 2>& det, std::vector<cv::Rect>& offset_box_vec,
		std::vector<float>& score_vec);

	void ScaleCoordinates(std::vector<Detection>& data, float pad_w, float pad_h,
		float scale, const cv::Size& img_shape);
	torch::Tensor xywh2xyxy(const torch::Tensor& x);

	std::vector<std::vector<Detection>> PostProcessing(const torch::Tensor& detections,
		float pad_w, float pad_h, float scale, const cv::Size& img_shape,
		float conf_thres, float iou_thres);

	std::vector<std::vector<Detection>> batchPostProcessing(const torch::Tensor& detections);
	torch::Tensor img2Tensor(cv::Mat img);
};
